import os
import yaml
import json # 仅用于测试输出

# 定义剧本配置文件的路径
PLAYBOOK_FILE_PATH = os.path.join(os.path.dirname(__file__), '..', 'config', 'playbooks.yaml')

# 加载剧本库
PLAYBOOK_LIBRARY = []
try:
    with open(PLAYBOOK_FILE_PATH, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
        if config and 'playbooks' in config:
            PLAYBOOK_LIBRARY = config['playbooks']
        else:
            print(f"警告: {PLAYBOOK_FILE_PATH} 文件格式不正确或缺少 'playbooks' 键。")
except FileNotFoundError:
    print(f"错误: 剧本配置文件 {PLAYBOOK_FILE_PATH} 未找到。请确保文件存在。")
except yaml.YAMLError as e:
    print(f"错误: 解析剧本配置文件 {PLAYBOOK_FILE_PATH} 失败: {e}")
except Exception as e:
    print(f"加载剧本配置文件时发生未知错误: {e}")

def find_playbooks_for_event(event_type: str) -> list:
    """
    根据事件类型查找匹配的剧本。
    """
    matched_playbooks = [p for p in PLAYBOOK_LIBRARY if p["event_type"].lower() == event_type.lower()]
    return matched_playbooks

if __name__ == "__main__":
    print("--- 测试 playbooks.py 中的剧本加载与查找 ---")

    if not PLAYBOOK_LIBRARY:
        print("剧本库为空，无法进行测试。请检查配置文件。")
    else:
        print(f"成功加载 {len(PLAYBOOK_LIBRARY)} 个剧本。")
        print("\n所有加载的剧本:")
        print(json.dumps(PLAYBOOK_LIBRARY, indent=2, ensure_ascii=False))

        print("\n查找 'malware_alert' 剧本:")
        playbooks = find_playbooks_for_event("malware_alert")
        print(json.dumps(playbooks, indent=2, ensure_ascii=False))

        print("\n查找 'reconnaissance_scan' 剧本:")
        playbooks = find_playbooks_for_event("reconnaissance_scan")
        print(json.dumps(playbooks, indent=2, ensure_ascii=False))

        print("\n查找 'unknown_event' 剧本:")
        playbooks = find_playbooks_for_event("unknown_event")
        print(json.dumps(playbooks, indent=2, ensure_ascii=False))